package org.example.nio.sendreceiver;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * @author hh
 * @since 2024/10/19
 */
public class NioReceiveServer {
    private static final String RECEIVE_PATH = "/Users/hh/test";
    private Charset charset = Charset.forName("UTF-8");

    private ByteBuffer buffer = ByteBuffer.allocate(1024);
    /**
     * 使用 Map保存每个客户端传输
     * 当 OP_READ通道可读时，根据Channel找到对应的对象
     */
    private Map<SelectableChannel, Client> clientMap = new HashMap<>();


    public void startServer() throws IOException {
        // 获取选择器
        Selector selector = Selector.open();
        // 获取通道
        ServerSocketChannel serverChannel = ServerSocketChannel.open();
        ServerSocket serverSocket = serverChannel.socket();
        // 设置为非阻塞
        serverChannel.configureBlocking(false);
        // 绑定链接
        InetSocketAddress address = new InetSocketAddress(9999);
        serverSocket.bind(address);
        // 将通道注册到选择器上，并且注册的 IO事件为"接收新连接"
        serverChannel.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("服务监听中...");
        // 轮询感兴趣的 IO就绪事件（选择键集合）
        while (selector.select() > 0) {
            // 获取选择键集合
            Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
            while (iterator.hasNext()) {
                // 获取单个的选择键并处理
                SelectionKey key = iterator.next();
                // 判断 key是具体的什么事件，是否为新连接事件
                if (key.isAcceptable()) {
                    // 若接收的事件是"新连接"，则获取客户端新连接
                    ServerSocketChannel server = (ServerSocketChannel) key.channel();
                    SocketChannel socketChannel = server.accept();
                    if (socketChannel == null) {
                        continue;
                    }
                    // 客户端新连接，切换为非阻塞模式
                    socketChannel.configureBlocking(false);
                    // 将客户端新连接通道注册到Selector上
                    socketChannel.register(selector, SelectionKey.OP_READ);
                    // 以下是业务处理
                    Client client = new Client();
                    client.remoteAddress = (InetSocketAddress) socketChannel.getRemoteAddress();
                    clientMap.put(socketChannel, client);
                    System.out.println(socketChannel.getRemoteAddress() + "连接成功");
                } else if (key.isReadable()) {
                    processData(key);
                }
                // nio的特点是只会累加，已选择键的集合不会删除，如果不删除，下一次又会被 select() 函数选中
                iterator.remove();

            }
        }
    }

    private void processData(SelectionKey key) throws IOException {
        Client client = clientMap.get(key.channel());
        SocketChannel socketChannel = (SocketChannel) key.channel();
        int num = 0;
        try {
            buffer.clear();
            while ((num = socketChannel.read(buffer)) > 0) {
                buffer.flip();
                // 客户端发送过来的，首先处理文件名
                if (null == client.fileName) {
                    if (buffer.capacity() < 4) {
                        continue;
                    }
                    int fileNameLen = buffer.getInt();
                    byte[] fileNameBytes = new byte[fileNameLen];
                    buffer.get(fileNameBytes);
                    // 文件名
                    String fileName = new String(fileNameBytes, charset);
                    File directory = new File(RECEIVE_PATH);
                    if (!directory.exists()) {
                        directory.mkdir();
                    }
                    System.out.println("NIO 传输目标dir: " + directory);
                    client.fileName = fileName;
                    String fullName = directory.getAbsolutePath() + File.separator + fileName;
                    System.out.println("NIO 传输目标文件：" + fullName);
                    File file = new File(fullName.trim());
                    if (!file.exists()) {
                        file.createNewFile();
                    }
                    client.outChannel = new FileOutputStream(file).getChannel();
                    if (buffer.capacity() < 8) {
                        continue;
                    }
                    // 文件长度
                    client.fileLength = buffer.getLong();
                    client.startTime = System.currentTimeMillis();
                    System.out.println("NIO 传输开始...");
                    client.receiveLength += buffer.capacity();
                    if (buffer.capacity() > 0 ){
                        // 写入文件
                        client.outChannel.write(buffer);
                    }
                    if (client.isFinished()) {
                        finished(key, client);
                    }
                    buffer.clear();
                }
                // 客户端发送过来的，最后是文件内容
                else {
                    client.receiveLength += buffer.capacity();
                    // 写入文件
                    client.outChannel.write(buffer);
                    if (client.isFinished()) {
                        finished(key, client);
                    }
                    buffer.clear();
                }
            }
            key.cancel();
        } catch (IOException e) {
            key.cancel();
            e.printStackTrace();
            return;
        }
        if (num == -1) {
            finished(key, client);
            buffer.clear();
        }
    }

    private void finished(SelectionKey key, Client client) throws IOException {
        client.outChannel.close();
        System.out.println("上传完毕...");
        key.cancel();
        System.out.println("文件接收成功，File name :" + client.fileName);
        System.out.println("size: "+ client.fileLength);
        long endTime = System.currentTimeMillis();
        System.out.println("传输时间（秒）：" + (endTime - client.startTime) / 1000);
    }


    class Client {
        /**
         * 文件名称
         */
        String fileName;
        /**
         * 长度
         */
        long fileLength;
        /**
         * 开始传输的时间
         */
        long startTime;
        /**
         * 客户端的地址
         */
        InetSocketAddress remoteAddress;

        /**
         * 输出的文件通道
         */
        FileChannel outChannel;

        /**
         * 接收的长度
         */
        long receiveLength;

        public boolean isFinished() {
            return receiveLength > fileLength;
        }
    }

}
